import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.loggers import TensorBoardLogger
import random
from argparse import ArgumentParser
from datasets.dataset import YOLODataset

def cli_main():
    # 设置随机种子
    seed = 522
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    # 以类的方式定义超参数
    parser = ArgumentParser()  # 创建参数解析器
    parser.add_argument('--backbone', default='resnet34d', type=str)  # 设置默认的 backbone（主干网络）
    parser.add_argument('--batch_size', default=48, type=int)  # 批量大小
    parser.add_argument('--epoch', default=50, type=int)  # 训练轮数
    parser.add_argument('--root', default='', type=str)  # 数据根目录
    parser.add_argument('--train_ann_path', default='', type=str)  # 训练集标注文件路径
    parser.add_argument('--num_class', default='1', type=int)  # 类别数
    parser.add_argument('--checkpoint', default='', type=str)  # 断点训练
    parser.add_argument('--pre_checkpoint', default='', type=str)  # 迭代训练
    parser.add_argument('--devices', default=[0, 1])  # 设置gpu
    args = parser.parse_args()  # 解析命令行参数

    # 载入数据
    train_dataset = YOLODataset(flag='train')
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
    valid_dataset = YOLODataset(flag='valid')
    valid_dataloader = DataLoader(dataset=valid_dataset, batch_size=64, shuffle=True)

    # 载入模型
    model = Your_model().to(args.device)
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(Your_model.parameters(), lr=args.learning_rate)

    # 模型训练




if __name__ == '__main__':
    cli_main()


